"""
生成模块
实现基于上下文的文本生成功能
"""

import os
from typing import List, Dict, Optional, Union
import logging
from dataclasses import dataclass
import json

from openai import AsyncOpenAI
import tiktoken

logger = logging.getLogger(__name__)

@dataclass
class GenerationResult:
    """生成结果"""
    answer: str
    sources: List[Dict]
    score: float
    
class Generator:
    """生成器类"""
    
    def __init__(self, config: Optional[Dict] = None):
        self.config = config or {}
        self.model_name = self.config.get("model_name", "gpt-3.5-turbo")
        self.temperature = self.config.get("temperature", 0.7)
        self.max_tokens = self.config.get("max_tokens", 1000)
        self.top_p = self.config.get("top_p", 1.0)
        
        # 初始化OpenAI客户端
        api_key = os.getenv("OPENAI_API_KEY")
        if not api_key:
            raise ValueError("未设置OPENAI_API_KEY环境变量")
            
        self.client = AsyncOpenAI(api_key=api_key)
        self.tokenizer = tiktoken.encoding_for_model(self.model_name)
        
        # 加载提示词模板
        self.prompt_template = self._load_prompt_template()
        
    def _load_prompt_template(self) -> str:
        """加载提示词模板"""
        return """你是一个专业的助手。请基于以下参考信息回答用户的问题。
如果参考信息不足以回答问题，请说明无法回答或信息不足。
请确保回答准确、客观，并标注信息来源。

用户问题：{question}

参考信息：
{context}

请按照以下格式回答：
1. 给出清晰、准确的回答
2. 说明使用了哪些参考信息
3. 如果信息不足，请说明原因

回答："""
        
    def _prepare_context(self, chunks: List[Dict]) -> str:
        """准备上下文信息"""
        context = []
        for i, chunk in enumerate(chunks, 1):
            doc_chunk = chunk["chunk"]
            source = f"[{i}] {doc_chunk.metadata.get('file_name', '未知来源')}"
            text = doc_chunk.text.strip()
            context.append(f"{source}:\n{text}\n")
        return "\n".join(context)
        
    def _count_tokens(self, text: str) -> int:
        """计算文本的token数量"""
        return len(self.tokenizer.encode(text))
        
    def _truncate_context(self, context: str, max_tokens: int) -> str:
        """截断上下文以适应token限制"""
        tokens = self.tokenizer.encode(context)
        if len(tokens) <= max_tokens:
            return context
            
        # 保留前面的tokens
        truncated_tokens = tokens[:max_tokens]
        return self.tokenizer.decode(truncated_tokens)
        
    async def generate(self, question: str, chunks: List[Dict]) -> GenerationResult:
        """生成回答"""
        try:
            # 准备上下文
            context = self._prepare_context(chunks)
            
            # 计算提示词的token数量
            prompt = self.prompt_template.format(
                question=question,
                context=context
            )
            
            # 如果提示词太长，截断上下文
            max_context_tokens = 4000 - self._count_tokens(question) - 500  # 保留空间给回答
            if self._count_tokens(context) > max_context_tokens:
                context = self._truncate_context(context, max_context_tokens)
                prompt = self.prompt_template.format(
                    question=question,
                    context=context
                )
                
            # 调用API生成回答
            response = await self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": "你是一个专业的助手，请基于给定的参考信息回答问题。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=self.temperature,
                max_tokens=self.max_tokens,
                top_p=self.top_p
            )
            
            # 解析回答
            answer = response.choices[0].message.content.strip()
            
            # 计算置信度分数
            score = sum(chunk["score"] for chunk in chunks) / len(chunks) if chunks else 0.0
            
            # 返回结果
            return GenerationResult(
                answer=answer,
                sources=chunks,
                score=score
            )
            
        except Exception as e:
            logger.error(f"生成回答失败: {str(e)}")
            raise
            
    async def generate_followup_questions(self, question: str, answer: str) -> List[str]:
        """生成后续问题"""
        try:
            prompt = f"""基于以下问答内容，生成3个相关的后续问题：

问题：{question}

回答：{answer}

请生成有助于深入理解主题的后续问题。每个问题应该探索不同的方面。
"""
            
            response = await self.client.chat.completions.create(
                model=self.model_name,
                messages=[
                    {"role": "system", "content": "你是一个专业的助手，请生成有见地的后续问题。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.8,
                max_tokens=200
            )
            
            # 解析回答，提取问题
            answer_text = response.choices[0].message.content.strip()
            questions = [q.strip() for q in answer_text.split("\n") if q.strip()]
            
            return questions[:3]  # 确保最多返回3个问题
            
        except Exception as e:
            logger.error(f"生成后续问题失败: {str(e)}")
            return [] 